package com.lzh.healthCode.utils;

import cn.hutool.core.util.HexUtil;
import cn.hutool.crypto.KeyUtil;
import cn.hutool.crypto.SmUtil;
import cn.hutool.crypto.asymmetric.KeyType;
import org.bouncycastle.crypto.engines.SM2Engine;

import java.nio.charset.StandardCharsets;
import java.security.KeyPair;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.util.HashMap;
import java.util.Map;
import java.util.function.Function;

/**
 * SM2 utility compatible with Java and C#
 *
 * @author Zhou Andong
 * @date 2022-7-1 16:14:11
 */
public interface SM2Util {

    String ALGORITHM = "SM2";

    String PUBLIC_KEY_NAME = "publicKey";

    String PRIVATE_KEY_NAME = "privateKey";

    Function<String, PublicKey> str2PublicKey = keyString -> KeyUtil
            .generatePublicKey(ALGORITHM, HexUtil.decodeHex(keyString));

    Function<String, PrivateKey> str2PrivateKey = keyString -> KeyUtil
            .generatePrivateKey(ALGORITHM, HexUtil.decodeHex(keyString));

    static Map<String, String> initKeyPair() {
        KeyPair keyPair = KeyUtil.generateKeyPair(ALGORITHM);
        Map<String, String> keyMap = new HashMap<>();
        keyMap.put(PUBLIC_KEY_NAME, HexUtil.encodeHexStr(keyPair.getPublic().getEncoded()));
        keyMap.put(PRIVATE_KEY_NAME, HexUtil.encodeHexStr(keyPair.getPrivate().getEncoded()));
        return keyMap;
    }

    static String encrypt(String publicKey, String data) {
        return SmUtil.sm2()
                .setMode(SM2Engine.Mode.C1C2C3)
                .setPublicKey(str2PublicKey.apply(publicKey))
                .encryptBase64(data, StandardCharsets.UTF_8, KeyType.PublicKey);
    }

    static String decrypt(String privateKey, String cipherData) {
        return SmUtil.sm2()
                .setMode(SM2Engine.Mode.C1C2C3)
                .setPrivateKey(str2PrivateKey.apply(privateKey))
                .decryptStr(cipherData, KeyType.PrivateKey);
    }

    // Demo
    public static void main(String[] args) {
        // 1. Generate key pair
        Map<String, String> keyPair = SM2Util.initKeyPair();
        String publicKey = keyPair.get("publicKey");
        String privateKey = keyPair.get("privateKey");
        Function<String, String> doEncrypt = data -> SM2Util.encrypt(publicKey, data);
        Function<String, String> doDecrypt = data -> SM2Util.decrypt(privateKey, data);

        // 2. Do encryption
        String data = "====Demo data to encrypt====";
        String encrypted = doEncrypt.apply(data);

        // 3. Do decryption
        String decrypted = doDecrypt.apply(encrypted);

        // 4. Show results
        System.out.println("publicKey: " + publicKey);
        System.out.println("privateKey: " + privateKey);

        System.out.println("--------------------");

        System.out.println("Original data: " + data);
        System.out.println("encrypted: " + encrypted);
        System.out.println("decrypted: " + decrypted);
    }

}
